"""
Name/Country 数据集,本例中的batch_size只能为1,因为name不等长
"""
import torch
from torch.utils.data import Dataset
import gzip
import csv


class Name2CountryDataset(Dataset):
    def __init__(self, isTrain: bool):
        """
        读取数据,根据输入的isTrain判定加载的数据集类型
        :param isTrain: True表示读取train的数据,False为test数据
        """
        super().__init__()
        # 根据是否为train训练集打开对应的文件
        filePath = "./dataset/names_train.csv.gz" if isTrain else "./dataset/names_test.csv.gz"
        # 打开文件
        with gzip.open(filePath, "rt") as f:
            # 使用csv读取
            reader = csv.reader(f)
            # 转为list,元素此时为(name,country)
            names = list(reader)
            # 分离list,获得单独的name的list和country的list
            rawNames, rawCountries = zip(*names)
            """处理name"""
            # 记录name映射成的ascii码数组
            self.names_ = []
            # 映射name
            for name in rawNames:
                # 单个name的ascii码数组
                tempSingleNameNum = []
                for c in name:
                    # 记录每个字符映射的数字
                    tempSingleNameNum.append(ord(c))
                tempSingleNameNum = torch.tensor(tempSingleNameNum, dtype=torch.long)
                # 记录
                self.names_.append(tempSingleNameNum)

            """处理country"""
            # country去重
            countries = list(set(rawCountries))
            # 排序
            countries.sort()
            # 记录country数目
            self.countrySize_ = len(countries)
            # 将country转换到字典
            self.idx2CountryMap_ = {idx: val for idx, val in enumerate(countries)}
            country2IdxMap = {val: idx for idx, val in enumerate(countries)}
            # 将国家映射到数字并记录
            self.countries_ = [country2IdxMap[x] for x in rawCountries]
            # 转换到tensor
            self.countries_ = torch.LongTensor(self.countries_)

    def __getitem__(self, index):
        return self.names_[index], self.countries_[index]

    def __len__(self):
        return len(self.names_)

    def countrySize(self):
        return self.countrySize_


if __name__ == '__main__':
    v = Name2CountryDataset(False)
